---
title: "Diagonstics of models"
output: html_notebook
---
# Part 4b: Diagnostics of models through traceplots

```{r}
library(rethinking)
library(bayesplot)
library(ggplot2)

label_map <- list(
  b_T_S = expression(b[t * "," * s]),
  b_C_S = expression(b[c * "," * s]),
  b_C_T = expression(b[c * "," * t]),
  b_C_W = expression(b[c * "," * w]),
  b_C_r = expression(b[c * "," * r]),
  b_T_W = expression(b[t * "," * w]),
  b_r_T = expression(b[r * "," * t]),
  b_r_W = expression(b[r * "," * w]),
  b_r_S = expression(b[r * "," * s])
)

trace_grid_theme <- theme(
  panel.grid.major = element_line(colour = "grey80", linewidth = 0.3),
  panel.grid.minor = element_line(colour = "grey90", linewidth = 0.2),
  panel.background = element_rect(fill = "white"),
  panel.border = element_rect(colour = "grey50", fill = NA, linewidth = 0.4)
)

make_y_lab <- function(param, model_num) {
  lhs <- label_map[[param]]
  if (is.expression(lhs)) lhs <- lhs[[1]]   # unwrap expression() to a single plotmath call

  bquote(.(lhs) ~ "(" * m[.(model_num)] * ")")
}

set_model_scheme <- function(model_key) {
  # bayesplot colour schemes (one family per model)
  scheme <- switch(
    model_key,
    m1 = "red",
    m2 = "orange",
    m3 = "green",
    m4 = "blue",
    "blue"
  )
  bayesplot::color_scheme_set(scheme)
}
```

# Load data and models from module 4
```{r}
# Load data from module_1
load("./data/d_sol_cw_w_in.RData")

# Load previous models
load("./data/model_out_final_real.RDATA")
load("./data/modeldata_out_final_real.RData")

dir.create("traceplots", showWarnings = FALSE)
```


# Traceplots for model m1
```{r}
set_model_scheme("m1")

# 1) Extract posterior draws from ulam
post <- rethinking::extract.samples(model_m1)  # a list of vectors/matrices

# 2) Convert to a matrix that bayesplot understands
draws <- posterior::as_draws_matrix(post)

# 3) Traceplots for only the parameters you want
for (p in c("b_C_S")) {

  g <- bayesplot::mcmc_trace(draws, pars = p) +
    labs(y = make_y_lab(p, model_num = 1)) +
    trace_grid_theme +
    theme(axis.title.y = element_text(size = 9))

  ggsave(
    filename = file.path("traceplots", paste0("trace_m1_", p, ".pdf")),
    plot = g,
    device = cairo_pdf,
    width = 60, height = 40,
    units = "mm"
  )
}
```

# Traceplots for model m2
```{r}
set_model_scheme("m2")

# 1) Extract posterior draws from ulam
post <- rethinking::extract.samples(model_m2)  # a list of vectors/matrices

# 2) Convert to a matrix that bayesplot understands
draws <- posterior::as_draws_matrix(post)

# 3) Traceplots for only the parameters you want
for (p in c("b_T_S", "b_C_T")) {

  g <- bayesplot::mcmc_trace(draws, pars = p) +
    labs(y = make_y_lab(p, model_num = 2)) +
    trace_grid_theme +
    theme(axis.title.y = element_text(size = 9))

  ggsave(
    filename = file.path("traceplots", paste0("trace_m2_", p, ".pdf")),
    plot = g,
    device = cairo_pdf,
    width = 60, height = 40,
    units = "mm"
  )
}
```

# Traceplots for model m3
```{r}
set_model_scheme("m3")

# 1) Extract posterior draws from ulam
post <- rethinking::extract.samples(model_m3)  # a list of vectors/matrices

# 2) Convert to a matrix that bayesplot understands
draws <- posterior::as_draws_matrix(post)

# 3) Traceplots for only the parameters you want
for (p in c("b_T_S", "b_C_S", "b_C_T")) {

  g <- bayesplot::mcmc_trace(draws, pars = p) +
    labs(y = make_y_lab(p, model_num = 3)) +
    trace_grid_theme +
    theme(axis.title.y = element_text(size = 9))

  ggsave(
    filename = file.path("traceplots", paste0("trace_m3_", p, ".pdf")),
    plot = g,
    device = cairo_pdf,
    width = 60, height = 40,
    units = "mm"
  )
}
```

# Traceplots for model m4
```{r}
set_model_scheme("m4")

# 1) Extract posterior draws from ulam
post <- rethinking::extract.samples(model_m4)  # a list of vectors/matrices

# 2) Convert to a matrix that bayesplot understands
draws <- posterior::as_draws_matrix(post)

# 3) Traceplots for only the parameters you want
for (p in c("b_T_S", "b_C_S", "b_C_T", "b_C_W", "b_C_r", "b_T_W", "b_r_T", "b_r_W", "b_r_S")) {

  g <- bayesplot::mcmc_trace(draws, pars = p) +
    labs(y = make_y_lab(p, model_num = 4)) +
    trace_grid_theme +
    theme(axis.title.y = element_text(size = 9))

  ggsave(
    filename = file.path("traceplots", paste0("trace_m4_", p, ".pdf")),
    plot = g,
    device = cairo_pdf,
    width = 60, height = 40,
    units = "mm"
  )
}
```